# scripts/train_ensemble.py
import os
import paddle
from models.efficient_time_llm import EfficientTimeLLM
from models.train_eval import train_model

def train_ensemble(X_train, time_train, y_train, save_dir='ensemble_models', num_models=5):
    """
    训练多个模型用于集成

    参数:
    - X_train: 训练历史数据 [N, seq_len, n_vars]
    - time_train: 时间特征 [N, seq_len, 2]
    - y_train: 预测目标 [N, pred_len]
    - save_dir: 模型保存路径
    - num_models: 集成模型数量
    """
    os.makedirs(save_dir, exist_ok=True)

    for i in range(num_models):
        print(f"\n开始训练第 {i + 1} 个模型")
        model = EfficientTimeLLM(seq_len=96, pred_len=24, d_model=128, num_layers=3, top_k=5, n_vars=X_train.shape[2])

        train_model(model, X_train, time_train, y_train, curriculum_steps=[6, 12, 24])

        paddle.save(model.state_dict(), f'{save_dir}/model_{i}.pdparams')
        print(f"✅ 第 {i + 1} 个模型训练完成并保存")

